from flask import Flask, render_template, request, jsonify
import numpy as np
import torch
import pickle
from model import TemperaturePredictor, load_model

app = Flask(__name__)

# 加载模型和相关组件
MODEL_PATH = 'temperature_model.pth'
SCALER_PATH = 'scaler.pkl'
FEATURE_LIST_PATH = 'feature_list.pkl'
INPUT_SIZE = 14  # 根据实际特征数量调整

# 加载模型
model, scaler, feature_list = load_model(
    MODEL_PATH, SCALER_PATH, FEATURE_LIST_PATH,
    input_size=INPUT_SIZE
)


@app.route('/')
def index():
    """渲染主页"""
    return render_template('index.html', feature_list=feature_list)


@app.route('/predict', methods=['POST'])
def predict():
    """处理预测请求"""
    try:
        # 获取表单数据
        data = request.form.to_dict()

        # 转换数值类型
        for key in data:
            if key != 'weekday':  # 星期几是字符串类型
                data[key] = float(data[key])

        # 处理星期几的独热编码
        weekday = data.pop('weekday')
        for day in ['Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat', 'Sun']:
            data[f'week_{day}'] = 1 if weekday == day else 0

        # 创建完整的输入字典
        input_dict = {}
        for feature in feature_list:
            # 确保所有特征都有值，缺失的特征设为0
            input_dict[feature] = data.get(feature, 0.0)

        # 进行预测
        prediction = predict_single_sample(model, scaler, feature_list, input_dict)

        return jsonify({
            'status': 'success',
            'prediction': round(prediction, 2)
        })

    except Exception as e:
        return jsonify({
            'status': 'error',
            'message': str(e)
        })


def predict_single_sample(model, scaler, feature_list, input_dict):
    """对单个样本进行预测"""
    # 1. 构造特征向量
    input_features = np.array([input_dict[feature] for feature in feature_list])

    # 2. 标准化
    input_features_scaled = scaler.transform(input_features.reshape(1, -1))

    # 3. 转换为Tensor
    input_tensor = torch.tensor(input_features_scaled, dtype=torch.float)

    # 4. 预测
    with torch.no_grad():
        prediction = model(input_tensor)

    return prediction.item()


if __name__ == '__main__':
    app.run(debug=True)